import json
import os
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
import torch
import copy
import numpy as np

from hand.detect.models.DetNet import DetNet
from hand.pipeline.detector.yvr_detector import DetNetDetector, YVRGTDetector
from hand.pipeline.loader.yvr_loader import YVRGTLoader
from hand.pipeline.tracker.dyd_tracker import DYDTrackerV2
from hand.pipeline.visualizer.hand_visualizer import draw_hand_keypoints
from hand.track.umetrack.datasets.video_pose_data import load_hand_model_from_dict
from hand.track.umetrack.hand.hand import HandModel, scaled_hand_model, mirrored_hand_model, NUM_JOINT_FRAMES, \
    hand_model_to_device
from hand.track.umetrack.hand.hand_pose_utils import landmarks_from_hand_pose, all_vertices_from_hand_pose
from hand.track.umetrack.hand.hand_skinning import skin_landmarks, _get_skinning_weights, _skin_points, \
    get_transformed_joint_positions
from hand.track.umetrack.hand.handpose import SingleHandPose
from hand.track.umetrack.model_v3.dyd_model import DYDModelV2
from hand.track.umetrack.model_v3.model_opts import ModelOpts

# 创建路由器
router = APIRouter(prefix="/model", tags=["model"])

# 全局变量
static_hand_model = None


# 请求模型
class UpdatePoseRequest(BaseModel):
    scale: float = 1.0
    joint_angles: list
    is_left_hand: bool = False


def load_hand_model():
    """加载手部模型"""
    global static_hand_model

    model_path = r"files/hand/models/track/generic_hand_model.json"
    try:
        with open(model_path, 'r') as f:
            model_data = json.load(f)

        # 转换数据为tensor
        joint_rotation_axes = torch.tensor(model_data['joint_rotation_axes'])
        joint_rest_positions = torch.tensor(model_data['joint_rest_positions'])
        joint_frame_index = torch.tensor(model_data['joint_frame_index'])
        joint_parent = torch.tensor(model_data['joint_parent'])
        joint_first_child = torch.tensor(model_data['joint_first_child'])
        joint_next_sibling = torch.tensor(model_data['joint_next_sibling'])
        landmark_rest_positions = torch.tensor(model_data['landmark_rest_positions'])
        landmark_rest_bone_weights = torch.tensor(model_data['landmark_rest_bone_weights'])
        landmark_rest_bone_indices = torch.tensor(model_data['landmark_rest_bone_indices'])
        hand_scale = torch.tensor(model_data['hand_scale'])
        mesh_vertices = torch.tensor(model_data['mesh_vertices'])
        mesh_triangles = torch.tensor(model_data['mesh_triangles'])
        dense_bone_weights = torch.tensor(model_data['dense_bone_weights'])
        joint_limits = torch.tensor(model_data['joint_limits'])

        static_hand_model = HandModel(
            joint_rotation_axes=joint_rotation_axes,
            joint_rest_positions=joint_rest_positions,
            joint_frame_index=joint_frame_index,
            joint_parent=joint_parent,
            joint_first_child=joint_first_child,
            joint_next_sibling=joint_next_sibling,
            landmark_rest_positions=landmark_rest_positions,
            landmark_rest_bone_weights=landmark_rest_bone_weights,
            landmark_rest_bone_indices=landmark_rest_bone_indices,
            hand_scale=hand_scale,
            mesh_vertices=mesh_vertices,
            mesh_triangles=mesh_triangles,
            dense_bone_weights=dense_bone_weights,
            joint_limits=joint_limits
        )
        return True
    except Exception as e:
        print(f"Error loading hand model: {str(e)}")
        return False


@router.get("/hand-model/")
async def get_hand_model():
    """获取手部模型数据"""
    try:
        if static_hand_model is None:
            if not load_hand_model():
                raise HTTPException(status_code=400, detail="Failed to load hand model")

        hand_model = static_hand_model

        # 转换为可序列化的格式，包含完整的模型数据
        model_data = {
            'joint_rest_positions': (hand_model.joint_rest_positions.numpy()/1000.0).tolist(),
            'landmark_rest_positions': (hand_model.landmark_rest_positions.numpy()/1000.0).tolist(),
            'joint_rotation_axes': hand_model.joint_rotation_axes.numpy().tolist(),
            'joint_frame_index': hand_model.joint_frame_index.numpy().tolist(),
            'joint_parent': hand_model.joint_parent.numpy().tolist(),
            'joint_first_child': hand_model.joint_first_child.numpy().tolist(),
            'joint_next_sibling': hand_model.joint_next_sibling.numpy().tolist(),
            'landmark_rest_bone_weights': hand_model.landmark_rest_bone_weights.numpy().tolist(),
            'landmark_rest_bone_indices': hand_model.landmark_rest_bone_indices.numpy().tolist(),
            'mesh_vertices': (hand_model.mesh_vertices.numpy()/1000.0).tolist() if hand_model.mesh_vertices is not None else None,
            'mesh_triangles': hand_model.mesh_triangles.numpy().tolist() if hand_model.mesh_triangles is not None else None,
            'dense_bone_weights': hand_model.dense_bone_weights.numpy().tolist() if hand_model.dense_bone_weights is not None else None,
            'joint_limits': hand_model.joint_limits.numpy().tolist() if hand_model.joint_limits is not None else None,
            'hand_scale': hand_model.hand_scale.numpy().tolist() if hand_model.hand_scale is not None else None,
        }

        # 验证数据格式
        assert isinstance(model_data['joint_rest_positions'], list), "joint_rest_positions should be a list"
        assert isinstance(model_data['joint_parent'], list), "joint_parent should be a list"
        assert len(model_data['joint_rest_positions']) == len(
            model_data['joint_parent']), "joint_rest_positions and joint_parent should have the same length"

        return model_data
    except Exception as e:
        print(f"Error in get_hand_model: {str(e)}")
        raise HTTPException(status_code=400, detail=str(e))


@router.post("/update-pose/")
async def update_hand_pose(request: UpdatePoseRequest):
    """更新手部姿态"""
    try:
        # 提取基础参数
        scale = request.scale
        joint_angles = torch.tensor(request.joint_angles)
        is_left_hand = request.is_left_hand

        if static_hand_model is None:
            if not load_hand_model():
                raise HTTPException(status_code=400, detail="Failed to load hand model")

        # 创建手部模型
        hand_model = copy.deepcopy(static_hand_model)

        if len(hand_model.joint_rotation_axes.shape) == 2:  # 如果还没有batch维度
            hand_model = HandModel(
                joint_rotation_axes=hand_model.joint_rotation_axes.unsqueeze(0),  # [22, 3] -> [1, 22, 3]
                joint_rest_positions=hand_model.joint_rest_positions.unsqueeze(0),  # [22, 3] -> [1, 22, 3]
                joint_frame_index=hand_model.joint_frame_index.unsqueeze(0),
                joint_parent=hand_model.joint_parent.unsqueeze(0),
                joint_first_child=hand_model.joint_first_child.unsqueeze(0),
                joint_next_sibling=hand_model.joint_next_sibling.unsqueeze(0),
                landmark_rest_positions=hand_model.landmark_rest_positions.unsqueeze(0),  # [21, 3] -> [1, 21, 3]
                landmark_rest_bone_weights=hand_model.landmark_rest_bone_weights.unsqueeze(0),  # [21, 3] -> [1, 21, 3]
                landmark_rest_bone_indices=hand_model.landmark_rest_bone_indices.unsqueeze(0),  # [21, 3] -> [1, 21, 3]
                hand_scale=hand_model.hand_scale.unsqueeze(0) if hand_model.hand_scale is not None else None,
                mesh_vertices=hand_model.mesh_vertices.unsqueeze(0) if hand_model.mesh_vertices is not None else None,
                # [788, 3] -> [1, 788, 3]
                mesh_triangles=hand_model.mesh_triangles.unsqueeze(0),
                dense_bone_weights=hand_model.dense_bone_weights.unsqueeze(
                    0) if hand_model.dense_bone_weights is not None else None,  # [788, 17] -> [1, 788, 17]
                joint_limits=hand_model.joint_limits.unsqueeze(0) if hand_model.joint_limits is not None else None
            )

        # 应用缩放
        hand_model = scaled_hand_model(hand_model, scale)

        # 应用镜像（如果需要）
        if not is_left_hand:
            hand_model = mirrored_hand_model(hand_model, to_mirror=torch.tensor([True]))

        # 创建手腕变换矩阵
        wrist_transform = torch.eye(4)

        try:
            # 计算新的手部姿态
            # 1. 更新关键点位置
            print("开始计算关键点位置...")
            # 确保所有输入张量具有相同的leading dimensions
            joint_angles = joint_angles.unsqueeze(0)  # [20] -> [1, 20]
            wrist_transform = wrist_transform.unsqueeze(0)  # [4, 4] -> [1, 4, 4]

            landmarks = skin_landmarks(
                hand_model,
                joint_angles,
                wrist_transform
            )
            landmarks /= 1000.0

            # 2. 更新网格顶点位置
            if hand_model.mesh_vertices is not None and hand_model.dense_bone_weights is not None:
                print("开始计算网格顶点位置...")
                # 使用网格顶点的骨骼权重进行蒙皮
                print("计算网格顶点蒙皮权重...")
                # 使用dense_bone_weights而不是landmark的权重
                skin_mat = _get_skinning_weights(
                    torch.arange(NUM_JOINT_FRAMES, device=hand_model.dense_bone_weights.device)
                    .unsqueeze(0)
                    .unsqueeze(0)
                    .expand(1, hand_model.dense_bone_weights.shape[1], -1),  # [1, 788, 17]
                    hand_model.dense_bone_weights,  # [1, 788, 17]
                    NUM_JOINT_FRAMES
                )

                # 计算网格顶点的蒙皮变换
                print("计算网格顶点变换...")
                mesh_vertices = _skin_points(
                    hand_model.joint_rest_positions,
                    hand_model.joint_rotation_axes,
                    skin_mat,
                    joint_angles,
                    hand_model.mesh_vertices,
                    wrist_transform
                )
                mesh_vertices /= 1000.0
            else:
                mesh_vertices = None
                print("跳过网格顶点更新：缺少必要的网格数据")

            # 3. 获取关节位置
            joint_positions = get_transformed_joint_positions(hand_model, joint_angles, wrist_transform)

        except Exception as e:
            print(f"计算过程中出错: {str(e)}")
            print(f"错误类型: {type(e)}")
            import traceback
            print(f"错误堆栈: {traceback.format_exc()}")
            raise

        # 返回完整的更新后的模型数据
        return {
            'landmarks': landmarks.squeeze(0).numpy().tolist(),
            'joint_positions': joint_positions.squeeze(0).numpy().tolist(),
            'mesh_vertices': mesh_vertices.squeeze(0).numpy().tolist() if mesh_vertices is not None else None,
            'mesh_triangles': hand_model.mesh_triangles.squeeze(
                0).numpy().tolist() if hand_model.mesh_triangles is not None else None,
            'joint_rest_positions': hand_model.joint_rest_positions.squeeze(0).numpy().tolist(),
            'joint_rotation_axes': hand_model.joint_rotation_axes.squeeze(0).numpy().tolist(),
            'joint_frame_index': hand_model.joint_frame_index.squeeze(0).numpy().tolist(),
            'joint_parent': hand_model.joint_parent.squeeze(0).numpy().tolist(),
            'joint_first_child': hand_model.joint_first_child.squeeze(0).numpy().tolist(),
            'joint_next_sibling': hand_model.joint_next_sibling.squeeze(0).numpy().tolist(),
            'landmark_rest_positions': hand_model.landmark_rest_positions.squeeze(0).numpy().tolist(),
            'landmark_rest_bone_weights': hand_model.landmark_rest_bone_weights.squeeze(0).numpy().tolist(),
            'landmark_rest_bone_indices': hand_model.landmark_rest_bone_indices.squeeze(0).numpy().tolist(),
            'hand_scale': hand_model.hand_scale.squeeze(
                0).numpy().tolist() if hand_model.hand_scale is not None else None,
            'dense_bone_weights': hand_model.dense_bone_weights.squeeze(
                0).numpy().tolist() if hand_model.dense_bone_weights is not None else None,
            'joint_limits': hand_model.joint_limits.squeeze(
                0).numpy().tolist() if hand_model.joint_limits is not None else None,
            'success': True
        }

    except Exception as e:
        print(f"Error in update_hand_pose: {str(e)}")
        print(f"错误类型: {type(e)}")
        import traceback
        print(f"错误堆栈: {traceback.format_exc()}")
        raise HTTPException(status_code=400, detail=str(e))


# 在您的后端 router 中添加以下接口

import asyncio
from fastapi import WebSocket, WebSocketDisconnect
import base64
import cv2
import numpy as np


class ConnectionManager:
    def __init__(self):
        self.active_connections = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_json(self, message: dict, websocket: WebSocket):
        await websocket.send_json(message)


manager = ConnectionManager()


@router.websocket("/ws/tracking")
async def websocket_tracking(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        # 初始化处理组件（与您的demo脚本类似）
        data_path = r"files/hand/datasets/yvr"
        loader = YVRGTLoader(data_path)

        # detect_model_path = r"C:\workspace\hand_xr\models\yvr_detect\detnet_best.pth"
        # detect_model = DetNet()
        # detect_model.eval()
        # detect_model.load_state_dict(torch.load(detect_model_path, map_location=torch.device('cpu')))
        # detector = DetNetDetector(loader, detect_model)
        detector = YVRGTDetector(loader)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model_path = r"files/hand/models/track/ckpt_epoch300.pt"
        weights = torch.load(model_path, map_location=device)["model_state_dict"]
        model_opts = ModelOpts()
        model = DYDModelV2(input_size=(96, 96), model_opts=model_opts)
        model.classification_head = None
        model.load_state_dict(weights)
        model.eval()

        hand_model_path = r"files/hand/models/track/generic_hand_model.json"
        with open(hand_model_path) as f:
            generic_hand_model = load_hand_model_from_dict(json.load(f))
            generic_hand_model = hand_model_to_device(generic_hand_model, device)
        tracker = DYDTrackerV2(model, device, generic_hand_model)

        common_hand_model = hand_model_to_device(loader.hand_model, device)

        count = 0
        new_tracking = {}
        use_history = False

        for input_frame, gt_tracking in loader:
            # Try to receive a control message from client, but don't block sending frames.
            # The frontend should send a 'start' after opening the socket; the server will
            # break when it receives 'stop'. Use a short timeout so the loop keeps running
            # even if the client doesn't send anything every frame.
            try:
                msg = await asyncio.wait_for(websocket.receive_text(), timeout=0.01)
                if msg == "stop":
                    break
            except asyncio.TimeoutError:
                # no message received this iteration, continue and send frame
                pass

            hand_samples = detector.detect(input_frame, gt_tracking)

            tracking_results = {0: None, 1: None}
            frame_data = {}

            for hand_idx, sample in hand_samples.items():
                if sample is None:
                    continue

                pred_joints, pred_wrist, skel_scales, pred_hand_idx = tracker.track(sample)
                hand_pose = SingleHandPose(
                    pred_joints.detach().numpy(),
                    pred_wrist.detach().numpy(),
                    np.array([1.0], dtype=np.float32)
                )
                scale_model = scaled_hand_model(common_hand_model, skel_scales)
                tracking_results[hand_idx] = {"pose": hand_pose, "model": scale_model}

                # 获取手部模型数据
                landmarks = landmarks_from_hand_pose(scale_model, hand_pose, hand_idx)
                mesh_vertices = all_vertices_from_hand_pose(scale_model, hand_pose, hand_idx)
                landmarks /= 1000.0
                mesh_vertices /= 1000.0
                # Ensure scale is a plain python float when serializing
                try:
                    scale_value = float(skel_scales.item()) if hasattr(skel_scales, 'item') else float(skel_scales)
                except Exception:
                    # fallback: try converting directly
                    scale_value = float(skel_scales)

                frame_data[f"hand_{hand_idx}"] = {
                    "landmarks": landmarks.tolist(),
                    "mesh_vertices": mesh_vertices.tolist(),
                    "joint_angles": hand_pose.joint_angles.tolist(),
                    "scale": scale_value,
                    "is_left_hand": hand_idx == 0  # 假设0是右手，1是左手
                }

            new_tracking = {}
            for hand_idx, sample in tracking_results.items():
                if sample is not None:
                    new_tracking[hand_idx] = landmarks_from_hand_pose(sample["model"], sample["pose"], hand_idx)

            # 处理视频帧
            processed_views = []
            for view in input_frame.views:
                view_image = view.image.copy()
                if len(view_image.shape) == 2:
                    view_image = cv2.cvtColor(view_image, cv2.COLOR_GRAY2BGR)
                elif view_image.shape[2] == 4:
                    view_image = cv2.cvtColor(view_image, cv2.COLOR_BGRA2BGR)

                # 绘制关键点
                for hand_idx, hand_result in tracking_results.items():
                    if hand_result is not None:
                        landmarks = landmarks_from_hand_pose(hand_result["model"], hand_result["pose"], hand_idx)
                        color = (0, 255, 0) if hand_idx == 0 else (255, 0, 0)
                        view_image = draw_hand_keypoints(view_image, landmarks, color, camera=view.camera)

                processed_views.append(view_image)
            # 合并视图（与您的visualizer类似）
            if len(processed_views) == 6:
                # 3x2 grid
                heights = [img.shape[0] for img in processed_views]
                widths = [img.shape[1] for img in processed_views]
                max_height = max(heights)
                max_width = max(widths)

                resized_views = []
                for img in processed_views:
                    if img.shape[0] != max_height or img.shape[1] != max_width:
                        img = cv2.resize(img, (max_width, max_height))
                    resized_views.append(img)

                top_row = np.hstack((resized_views[0], resized_views[1]))
                middle_row = np.hstack((resized_views[2], resized_views[3]))
                bottom_row = np.hstack((resized_views[4], resized_views[5]))
                grid_image = np.vstack((top_row, middle_row, bottom_row))
                grid_image = cv2.resize(grid_image, (max_width, int(max_height * 1.5)))
            else:
                grid_image = processed_views[0]  # 默认使用第一个视图

            # 编码图像为base64
            _, buffer = cv2.imencode('.jpg', grid_image)
            frame_base64 = base64.b64encode(buffer).decode('utf-8')

            camera = input_frame.views[0].camera
            camera_to_world_xf = camera.camera_to_world_xf
            extrinsic = np.linalg.inv(camera_to_world_xf)
            # extrinsic = camera_to_world_xf

            f = camera.f
            c = camera.c
            w = camera.width
            h = camera.height

            # 发送数据到前端
            await manager.send_json({
                "type": "frame",
                "frame_index": count,
                "video_frame": f"data:image/jpeg;base64,{frame_base64}",
                "hand_data": frame_data,
                "extrinsic": extrinsic.tolist(),  # 将numpy数组转换为列表
                "timestamp": count / 30.0,  # 假设30fps,
                "f": f,
                "c": c,
                "w": w,
                "h": h
            }, websocket)

            count += 1
            # if count > 1000:  # 限制帧数，避免无限循环
            #     break

    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        print(f"WebSocket error: {str(e)}")
        manager.disconnect(websocket)


# 初始化时加载手部模型
def init_hand_model():
    """初始化手部模型"""
    return load_hand_model()


# 应用启动时加载手部模型
init_hand_model()
